# -*- coding: utf-8 -*-
# created on 2016/5/5

from sympy import sympify, Eq, FiniteSet, Union, EmptySet, Set, Symbol, S, Expr
import itertools
from mathsolver.functions.sympy_utils import ReplaceAll
from mathsolver.functions.base.ordered_set import OrderedSet

latex_symbol = {"belong": "\\in", "because": "∵"}


def search_multiple(known, names):  # todo: one key save all values !
    select_key = [name for name in known if
                  isinstance(name, tuple) and 0 < len(set(names).intersection(name)) <= len(names)]
    if not select_key:
        return

    select_key = select_key[0]
    names = [n for n in names if n in select_key]  # list(set(select_key).intersection(names))
    temp = [str(k) for k in select_key]
    select_index = [temp.index(str(n)) for n in names]
    cond_index = [i for i, n in enumerate(select_key) if
                  n not in names and any([isinstance(v[i], Set) for v in known[select_key]])]
    piecewise = PieceWise(names)
    for value in known[select_key]:
        select_value = [
            select_key[index] if isinstance(value[index], Set) and value[index].free_symbols else value[index] for index
            in select_index]
        cond = {}
        for i, v in enumerate(value):
            if i not in select_index and i in cond_index:
                cond[select_key[i]] = v

            if i in select_index and isinstance(v, Set):
                cond[select_key[i]] = v

        if len(select_value) > 1:
            piecewise.add(tuple(select_value), cond)
        else:
            piecewise.add(select_value[0], cond)

    return piecewise.simplify()


def print_ans(ans):
    keys, values = list(ans.items())[0]
    if isinstance(keys, Symbol):
        if isinstance(values, FiniteSet):
            return "%s = %s" % (new_latex(keys), ",".join([new_latex(v) for v in values]))
        else:
            return "%s %s %s" % (new_latex(keys), latex_symbol["belong"], new_latex(values))
    else:
        latex_exprs = []
        for value in values:
            latex_expr = []
            for k, v in zip(keys, value):
                if isinstance(v, Set):
                    latex_expr.append(new_latex(k) + latex_symbol["belong"] + new_latex(v))
                else:
                    latex_expr.append(new_latex(k) + "=" + new_latex(v))

            latex_exprs.append(",".join(latex_expr))

        if len(latex_exprs) == 1:
            return latex_exprs[0]
        else:
            latex_exprs = "\\\\".join(latex_exprs)
            return "\\begin{equation}\\begin{cases}%s\\end{cases}\\end{equation}" % latex_exprs


class BaseValue(object):
    intermediate = False
    is_or = False
    main = ""

    def __init__(self, value=None):
        self.value = value
        self.text = self.__class__.__name__
        self.latex = None

    def sympify(self):
        """ 输出sympy格式 """
        if isinstance(self, Expr):
            return self

        try:
            return sympify(self.value, evaluate=False)
        except Exception:
            return sympify(self.value)

    def printing(self):
        """ 输出latex格式 """
        try:
            return new_latex(self.sympify())
        except Exception:
            return str(self.value)

    def to_query(self):
        return self.__class__.__name__.replace("Base", "")

    def get_input(self, name, exprs):
        if "func" in name:
            assert isinstance(self.value, list) and self.value[0] == "y"
        return self

    def free_symbols(self):
        return set()

    def add_latex(self, value):
        self.latex = value


class BaseFunction(object):
    def __init__(self, known=None, verbose=False, target=None, exist_symbols=None):
        if target is None:
            target = []
        if known is None:
            known = {}
        self.steps = []
        self.output = []
        self.verbose = verbose
        self.label = OrderedSet()
        self.new_conditions = []
        self.new_questions = []
        self.known = known
        self.target = target
        self.exist_symbols = exist_symbols

    def search(self, name, known=None):
        if not known:
            known = self.known
        if name == "NN" and name not in known:
            name = "N"

        if name == "N" and name not in known:
            name = "NN"

        if name in known:
            return known[name]

        elif name == "O":
            return BaseValue([0, 0])
        else:
            for key, value in known.items():
                if isinstance(key, tuple) and name in key:
                    index = key.index(name)
                    value = [v[index] for v in value]
                    if len(value) == 1 and isinstance(value[0], Set):
                        return value[0]
                    else:
                        return FiniteSet(*value)
                elif str(key) == str(name):
                    return value
                elif str(key).find(str(name)) >= 0:  # search a_S_类
                    import re
                    p = r'[A-Z]_'
                    m = re.findall(p, str(key))
                    if m:
                        return value

            return None

    def searchs(self, names):
        values = {}
        if len(names) > 1:
            for name in itertools.permutations(names, len(names)):
                name = tuple(name)
                if name in self.known:
                    return {name: self.known[name]}

        for name in names:
            value = self.search(name)
            if value:
                values[name] = value

        return values

    def search_multiple(self, names):  # todo: one key save all values !
        return search_multiple(self.known, names)

    def solver(self, *args):
        return self

    def append(self, stepsolver, output=True):
        self.steps += stepsolver.steps
        if output:
            self.output += stepsolver.output
        self.label.update(stepsolver.label)

    def printing(self):
        steps = []
        for a, b in self.steps:
            a = a.replace("$", "")
            b = b.replace("$", "")
            steps.append([a, b])

        return steps

    @staticmethod
    def output_jie(output, symbol, eq="="):
        return ",".join(["%s %s %s" % (new_latex(symbol), eq, new_latex(out)) for out in output])

    @staticmethod
    def output_jies(output):
        known_sym, known_solve = zip(*output.items())
        answers = []
        for known in zip(*known_solve):
            known = [new_latex(Eq(s, a)) for s, a in zip(known_sym, known)]
            answer = "\\\\".join(known)
            answers.append("\\begin{equation}\\begin{cases}%s\\end{cases}\\end{equation}" % answer)
        return ",".join(answers)

    @staticmethod
    def output_symbol_value(symbol_value):
        return ",".join(["%s = %s" % (new_latex(symbol), new_latex(value)) for symbol, value in symbol_value.items()])

    @staticmethod
    def _output_interval(symbol, value):
        if isinstance(value, FiniteSet):
            return "%s = %s" % (new_latex(symbol), ",".join([new_latex(n) for n in value]))
        l, r, lopen, ropen = value.args
        smaller = {False: "<=", True: "<"}
        larger = {False: ">=", True: ">"}
        if l.is_infinite:
            if r.is_infinite:  # (负无穷大，正无穷大)
                return "%s \\in R" % new_latex(symbol)
            else:  # (负无穷大，m)
                return "%s %s %s" % (new_latex(symbol), smaller[ropen], new_latex(r))
        else:
            if r.is_infinite:  # (m，正无穷大)
                return "%s %s %s" % (new_latex(symbol), larger[lopen], new_latex(l))
            else:  # (n，m)
                return "%s %s %s %s %s" % (new_latex(l), smaller[lopen], new_latex(symbol), smaller[ropen], new_latex(r))

    def output_interval(self, symbol, value):
        if isinstance(value, EmptySet):
            return "无解"
        elif isinstance(value, Union):
            return " 或 ".join([self._output_interval(symbol, arg) for arg in value.args])
        else:
            return self._output_interval(symbol, value)

    def output_answer(self, symbol_value):
        if len(symbol_value) > 1:
            return self.output_jies(symbol_value)
        else:
            symbol, value = list(symbol_value.items())[0]
            if isinstance(value, list) or isinstance(value, FiniteSet):
                return ",".join(["%s = %s" % (new_latex(symbol), new_latex(v)) for v in value])
            else:
                return self.output_interval(symbol, value)

    @staticmethod
    def output_eq(eq):
        if len(eq) == 3:
            f1, op, f2 = eq
            return "%s %s %s" % (new_latex(f1), op, new_latex(f2))
        else:
            f1, f2 = eq
            return "%s = %s" % (new_latex(f1), new_latex(f2))

    def output_eqs(self, eqs):
        latex_eqs = []
        for eq in eqs:
            latex_eqs.append(self.output_eq(eq))

        latex_eqs = "\\\\".join(latex_eqs)
        return "\\begin{equation}\\begin{cases}%s\\end{cases}\\end{equation}" % latex_eqs


class PieceWise(object):
    def __init__(self, variables):
        self.pieces = []
        self.var = variables

    def add(self, value, cond):
        self.pieces.append([value, cond])

    def simplify(self):
        # if all([not cond for value, cond in self.pieces]):
        #     return FiniteSet(*[value for value,_ in self.pieces])
        # else:
        pieces = sorted(self.pieces, key=lambda x: str(x[0]))
        new_pieces = []
        for group, items in itertools.groupby(pieces, lambda x: str(x[0])):
            items = list(items)
            value = items[0][0]
            conds = [cond for _, cond in items]
            merged_cond = conds[0]
            for cond in conds[1:]:
                merged_cond = self.merge(merged_cond, cond)

            merged_cond = dict([(k, v) for k, v in merged_cond.items() if v != S.Reals])
            new_pieces.append([value, merged_cond])

        self.pieces = new_pieces
        # if all([not cond for value, cond in self.pieces]):
        #     return FiniteSet(*[value for value,_ in self.pieces])
        return self

    def printing(self):
        latex_exprs = []
        for value, cond in self.pieces:
            latex_exprs.append("%s,%s" % (new_latex(value),
                                          ",".join([new_latex(k) + latex_symbol["belong"] + new_latex(v)
                                                    if isinstance(v, Set) else new_latex(k) + "=" + new_latex(v) for
                                                    k, v in cond.items()])))

        latex_exprs = "\\\\".join(latex_exprs)
        return "\\begin{equation}\\begin{cases}%s\\end{cases}\\end{equation}" % latex_exprs

    @staticmethod
    def merge(merged_cond, cond):
        assert set(merged_cond.keys()) == set(cond.keys())
        new_cond = {}
        for k, v in merged_cond.items():
            if not isinstance(v, Set):
                v = FiniteSet(v)
            new_v = FiniteSet(cond[k]) if not isinstance(cond[k], Set) else cond[k]
            new_value = Union(v, new_v)
            new_cond[k] = new_value

        return new_cond


def new_latex(expr):
    from sympy import latex
    if isinstance(expr, Expr):
        if "Angle" in str(expr) or "distance" in str(expr) or "vector" in str(expr) or "_" in str(
                expr) or "binomial" in str(expr):
            replacer = ReplaceAll("aaa")
            new_expr = replacer.replace(expr, lambda x: hasattr(x, "func") and (
                        str(x.func) in ["Angle", "distance", "vector", "binomial"] or str(x.func)[-1] == "_"))
            expr = latex(new_expr)
            for key, value in replacer.symbol_value.items():
                if str(value.func) == "Angle":
                    value = "\\angle " + "".join([latex(arg) for arg in value.args])
                elif str(value.func) == "vector":
                    value = "\\overrightarrow{%s}" % "".join([latex(arg) for arg in value.args])
                elif str(value.func) == "distance":
                    value = "|%s|" % "".join([latex(arg) for arg in value.args])
                elif str(value.func) == "binomial":
                    value = "C_{%s}^%s" % (latex(value.args[0]), latex(value.args[1]))
                elif str(value.func)[-1] == "_":
                    value = "%s{%s}" % (str(value.func), latex(value.args[0]))
                expr = expr.replace(latex(key), value)
        else:
            expr = latex(expr)
    else:
        expr = latex(expr)

    expr = expr.replace("\\left", "").replace("\\right", "").replace("\\frac", "\\dfrac")
    expr = expr.replace("beta_", "\\beta").replace("lambda_", "\\lambda ").replace("gamma_", "\\gamma ")
    return expr


if __name__ == '__main__':
    pass
